#define _GNU_SOURCE

#include <link.h>
#include <elf.h>
#include <stdlib.h>
#include <string.h>
#include <sys/wait.h>
#include <sys/types.h>
#include <stdio.h>
#include <errno.h>
#include <unistd.h>

#include "map_utils_linux_elf.h"
#include "linux_elf_utils.h"
#include "map_utils.h"
#include "memory_utils.h"
#include "debug.h"

//------------------------------------------------------------------------------

TBGError ElfTryGetFuncMap(elf_session *pelf, func_map *pmap) {
    TBGError ret = 0;
    CHECK_DBG_ERR(pmap, NULL_PTR);
    CHECK_DBG_ERR(pelf, NULL_PTR);

    int cur_size = SYMBOLS_COUNT_START;
    int read_count = 0;
    pmap->map_ = (func_map_info*) malloc(cur_size * sizeof(func_map_info));
    pmap->names_ = (name_t*) malloc(cur_size * sizeof(name_t));
    CHECK_RET(pmap->map_, CANNOT_ALLOC_MEMORY, free_point, ret);
    CHECK_RET(pmap->names_, CANNOT_ALLOC_MEMORY, free_point, ret);
 

    if(pelf->symtab && pelf->nchains > 0) //may be old symbols
        for(;;) { 
            CHECK_CALL(ElfGetAllSymbolsOld(pelf, pmap->map_, pmap->names_,
                                           cur_size, &read_count),
                       free_point, ret);
            if(read_count < cur_size) // we have read all dynamic symbols
                break;
            cur_size *= 2;

            free(pmap->map_);
            free(pmap->names_);
            pmap->map_ = (func_map_info*) malloc(cur_size *
                                                 sizeof(func_map_info));
            pmap->names_ = (name_t*) malloc(cur_size * sizeof(name_t));
            CHECK_RET(pmap->map_, CANNOT_ALLOC_MEMORY, free_point, ret);
            CHECK_RET(pmap->names_, CANNOT_ALLOC_MEMORY, free_point, ret);
        } 

    int read_old = read_count;
    read_count = 0;

    for(;;) { 
        CHECK_CALL(ElfGetAllSymbolsGNU(pelf, pmap->map_ + read_old,
                                       pmap->names_, cur_size - read_old,
                                       &read_count),
                   free_point, ret);
        if(read_count < cur_size - read_old) // we have read all dynamic symbols
            break;
        cur_size *= 2;

        pmap->map_ = (func_map_info*) realloc(pmap->map_,
                                              cur_size * sizeof(func_map_info));
        pmap->names_ = (name_t*) realloc(pmap->names_,
                                         cur_size * sizeof(name_t));
        CHECK_RET(pmap->map_, CANNOT_ALLOC_MEMORY, free_point, ret);
        CHECK_RET(pmap->names_, CANNOT_ALLOC_MEMORY, free_point, ret);
    } 

    pmap->size_ = read_old + read_count;

    DEBUG("pmap size is [%x]", pmap->size_);


exit_point:
    DEBUG("return [%s]", StrError(ret));
    return ret;
free_point:
    free(pmap->map_);
    free(pmap->names_);
    goto exit_point;
}

//------------------------------------------------------------------------------

TBGError ElfGetDynamicFuncMap(func_map *pmap) {
    TBGError ret = 0;
    CHECK_DBG_ERR(pmap, NULL_PTR);

    func_map temp_map = {};
    struct link_map l_map = {};
    elf_session elf = {};
    ElfLocateLinkmap(ELF_HEADER_BEGIN, &l_map);

    for(;;){  
        temp_map.size_ = 0;
        temp_map.map_ = temp_map.names_ = NULL;

        CHECK_CALL(ElfResolveTables(&l_map, &elf), exit_point, ret); 

        CHECK_CALL(ElfTryGetFuncMap(&elf, &temp_map), exit_point, ret); 

        CHECK_CALL(JoinToFuncMap(pmap, &temp_map), free_point, ret); 

        free(temp_map.map_);
        free(temp_map.names_);
        
        if(!l_map.l_next)
            break;

        CHECK_CALL(ReadByteData(l_map.l_next, &l_map, sizeof(struct link_map)),
                   exit_point, ret);
    }

exit_point:
    return ret;
free_point:
    free(temp_map.map_);
    free(temp_map.names_);
    goto exit_point;
}

//------------------------------------------------------------------------------

TBGError ElfResolveTables(struct link_map *pl_map, elf_session *pelf) {
    TBGError ret = 0;
    CHECK_DBG_ERR(pelf, NULL_PTR);

	Elf32_Dyn dyn = {};

    DEBUG("l_map.l_addr = [%x]", pl_map->l_addr);

    CHECK_CALL(ElfGetDynSym(pl_map, DT_HASH, &dyn), strtab_got, ret);

    DEBUG("dyn.d_un_ptr = [%x], l_map.l_addr = [%x], addr is [%x]",
          dyn.d_un.d_ptr, pl_map->l_addr, dyn.d_un.d_ptr + pl_map->l_addr);

    CHECK_CALL(ReadByteData(dyn.d_un.d_ptr + /*pl_map->l_addri*/ +
                            sizeof(Elf32_Dyn), &(pelf->nchains), 
                            sizeof(pelf->nchains)), strtab_got, ret);
strtab_got:

    CHECK_CALL(ElfGetDynSym(pl_map, DT_STRTAB, &dyn), symtab_got, ret);
    pelf->strtab = dyn.d_un.d_ptr;

symtab_got:

    CHECK_CALL(ElfGetDynSym(pl_map, DT_SYMTAB, &dyn), gnu_got, ret);
    pelf->symtab = dyn.d_un.d_ptr;

gnu_got:

    CHECK_CALL(ElfGetDynSym(pl_map, DT_GNU_HASH, &dyn), exit_point, ret);
    pelf->gnutab = dyn.d_un.d_ptr;

exit_point:
    return ret == DYN_SYM_NOT_FOUND ? 0 : ret;
}

//------------------------------------------------------------------------------

int ElfGetMaxOldSymbolsCount(elf_session *pelf) { 
    CHECK_DBG_ERR(pelf, NULL_PTR);
    CHECK_DBG_ERR(pelf->nchains, INVALID_CHAINS_COUNT);
    return pelf->nchains;
}

//------------------------------------------------------------------------------

TBGError ElfGetAllSymbolsOld(elf_session *pelf, func_map_info *map_arr_out,
                             name_t *names, size_t max_size, int *got_count) {
    TBGError ret = 0;
    CHECK_DBG_ERR(pelf, NULL_PTR);
    CHECK_DBG_ERR(map_arr_out, NULL_PTR);
    CHECK_DBG_ERR(names, NULL_PTR);
    CHECK_DBG_ERR(max_size, ZERO_LENGTH);
    if(pelf->nchains) {
        CHECK_RET(pelf->symtab, INVALID_SYMTAB_ADDR, exit_point, ret);
        CHECK_RET(pelf->strtab, INVALID_STRTAB_ADDR, exit_point, ret);
    }

	Elf32_Sym sym = {};
    int count = 0;
	int i = 0;
	for (i = 0; i < pelf->nchains && i < max_size; ++i) {
		CHECK_CALL(ReadByteData(pelf->symtab + (i * sizeof(Elf32_Sym)), &sym,
                                sizeof(Elf32_Sym)), 
                   exit_point, ret);
		if (ELF32_ST_TYPE(sym.st_info) != STT_FUNC)
			continue;
        // our symbol is a function name
        map_arr_out[count].addr_ = sym.st_value;
        map_arr_out[count].name_idx_ = count;

        // read identifier name from string table (STR_TAB)
		CHECK_CALL(ReadStr(pelf->strtab + sym.st_name, func_name_length_,
                           names[count]),
                   exit_point, ret);
        
        count++;
	}

    *got_count = count;

exit_point:
    return ret;
}

//------------------------------------------------------------------------------

/*static uint_fast32_t dl_new_hash (const char *s) {
    uint_fast32_t h = 5381;
    unsigned char c = 0;
    for (c = *s; c != '\0'; c = *++s)
        h = h * 33 + c;
    return h & 0xffffffff;
}*/

//------------------------------------------------------------------------------

TBGError ElfGetAllSymbolsGNU(elf_session *pelf, func_map_info *map_arr_out, 
                             name_t *names, size_t max_size, int *got_count) {
    TBGError ret = 0;
    CHECK_DBG_ERR(pelf, NULL_PTR);
    CHECK_DBG_ERR(map_arr_out, NULL_PTR);
    CHECK_DBG_ERR(names, NULL_PTR);
    CHECK_DBG_ERR(max_size, ZERO_LENGTH);

	Elf32_Dyn dyn = {};

    CHECK_RET(pelf->gnutab, INVALID_HASHMAP_ADDR, exit_point, ret);
    CHECK_RET(pelf->symtab, INVALID_SYMTAB_ADDR, exit_point, ret);
    CHECK_RET(pelf->strtab, INVALID_STRTAB_ADDR, exit_point, ret);

    Elf32_Word nbuckets = 0;//pelf->gnutab[0];
    Elf32_Word symbndx  = 0;//pelf->gnutab[1];
    Elf32_Word nmaskwords = 0;//pelf->gnutab[2];

    DEBUG("pelf->gnutab is [%x]", pelf->gnutab);

    CHECK_CALL(ReadByteData(pelf->gnutab + 0 * sizeof(Elf32_Word), &nbuckets,
                            sizeof(nbuckets)),
               exit_point, ret);
    CHECK_CALL(ReadByteData(pelf->gnutab + 1 * sizeof(Elf32_Word), &symbndx, 
                            sizeof(symbndx)),
               exit_point, ret);
    CHECK_CALL(ReadByteData(pelf->gnutab + 2 * sizeof(Elf32_Word), &nmaskwords,
                            sizeof(nmaskwords)),
               exit_point, ret);

    DEBUG("nbuckets is [%x]", nbuckets);

    //pelf->gnutab + 3 is shift param to bloom filter. Ignore it.
    //second section is bloom filter bitmask.
    unsigned long bitmask = pelf->gnutab + 4 * sizeof(Elf32_Word);

    //we should skip bloom filter mask to get the fird section
    unsigned long buckets = bitmask + nmaskwords * sizeof(ElfW(Addr));
    //and skip the buckets section to get pointer to fourth hashval section
    unsigned long hashvals  = buckets +
                              (nbuckets - symbndx) * sizeof(Elf32_Word);

    Elf_Symndx idx = 0;
    Elf32_Sym sym;// = {};
    //we need to get all

    size_t count = 0;
    Elf32_Word i = 0;
    Elf32_Word next_hash = 0;
    for(i = 0; i < nbuckets; ++i) {
        CHECK_CALL(ReadByteData(buckets + i * sizeof(Elf32_Word), &idx,
                                sizeof(idx)), 
                   exit_point, ret);

        if(!idx) // there is no chain with this hash
            continue;

        do {
            // go throw the chain
            CHECK_CALL(ReadByteData(pelf->symtab + idx * sizeof(Elf32_Sym),
                                    &sym, sizeof(Elf32_Sym)),
                       exit_point, ret);
            
            if (ELF32_ST_TYPE(sym.st_info) == STT_FUNC) {
                CHECK_DBG_ERR(count < max_size, BIG_INDEX);

                map_arr_out[count].addr_ = sym.st_value;
                map_arr_out[count].name_idx_ = count;

                // read identifier name from string table (STR_TAB)
                CHECK_CALL(ReadStr(pelf->strtab + sym.st_name,
                                   func_name_length_, names[count]),
                           exit_point, ret);
                
                count++;
                if(count == max_size) 
                    goto exit_point;
            }

            CHECK_CALL(ReadByteData(hashvals + idx * sizeof(Elf32_Word),
                                    &next_hash, sizeof(Elf32_Word)),
                       exit_point, ret);
            idx++;

        } while((next_hash & 1u) != 1u); // last bit of hash is the end marker
    }

exit_point:
    *got_count = count;
    return ret;
}

